﻿#include <iostream>
#include <stack>
#include "UnionFind.h"

template<class T>
UnionFind<T>::UnionFind(vector<T> list)
{
	for (auto itor = list.begin(); itor != list.end(); ++itor)
	{
		auto element = new Element<T>(*itor);
		_elements.insert({ *itor, element });
		_elementsToHeads.insert({ element, element });
		_elementSizes.insert({ element, 1 });
	}
}

template<class T>
UnionFind<T>::~UnionFind()
{
	_elementSizes.clear();
	_elementsToHeads.clear();
	for (auto itor = _elements.begin(); itor != _elements.end(); ++itor)
	{
		auto element = itor->second;
		delete(element);
	}

	_elements.clear();
}

template<class T>
Element<T>* UnionFind<T>::findHead(Element<T>* element)
{
	stack<Element<T>*> elementStack;
	elementStack.push(element);
	Element<T>* parent = nullptr;
	while (!elementStack.empty())
	{
		parent = _elementsToHeads.at(element);
		if (parent == element) break;

		elementStack.push(parent);
		element = parent;
	}

	while (!elementStack.empty())
	{
		auto top = elementStack.top();
		_elementsToHeads.at(top) = element;
		elementStack.pop();
	}

	return element;
}

template<class T>
void UnionFind<T>::unionElements(T t1, T t2)
{
	auto element1 = _elements.find(t1);
	auto element2 = _elements.find(t2);
	if (element1 == _elements.end() || element2 == _elements.end()) return;

	auto head1 = findHead(_elements.at(t1));
	auto head2 = findHead(_elements.at(t2));
	if (head1 == head2) return;

	size_t size1 = _elementSizes.at(head1);
	size_t size2 = _elementSizes.at(head2);
	auto parent = size1 > size2 ? head1 : head2;
	auto child = size1 > size2 ? head2 : head1;
	_elementsToHeads.at(child) = parent;
	_elementSizes.at(parent) = size1 + size2;
	_elementSizes.erase(child);
}

template<class T>
bool UnionFind<T>::isTheSameHead(T t1, T t2)
{
	auto element1 = _elements.find(t1);
	auto element2 = _elements.find(t2);
	if (element1 == _elements.end() || element2 == _elements.end()) return false;

	auto head1 = findHead(_elements.at(t1));
	auto head2 = findHead(_elements.at(t2));
	return head1 == head2;
}

// 并查集
int main_UnionFind()
{
	vector<int> nums;
	nums.push_back(1);
	nums.push_back(2);
	nums.push_back(3);
	nums.push_back(4);
	nums.push_back(5);

	UnionFind<int>* unionFind = new UnionFind<int>(nums);
	bool isSame = unionFind->isTheSameHead(1, 2);
	unionFind->unionElements(1, 2);
	isSame = unionFind->isTheSameHead(1, 2);
	isSame = unionFind->isTheSameHead(3, 2);
	unionFind->unionElements(3, 2);
	isSame = unionFind->isTheSameHead(3, 2);

	delete(unionFind);
	return 0;
}